from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.primitives.asymmetric import rsa, padding as asym_padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.backends import default_backend
import hashlib
import os
import base64

class CryptoUtils:

    def __init__(self):
        self.backend = default_backend()

    # AES 加密
    def aes_encrypt(self, data, key, iv):
        """
        使用AES加密算法对输入的字符串进行加密
        
        Args:
            data: 需要加密的字符串
            key: AES加密的密钥，长度为16、24或32字节
            iv: AES加密的初始化向量，长度为16字节
        
        Returns:
            加密后的字符串，使用base64编码
        
        """
        data = data.encode('utf-8')  # 将字符串转换为字节
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(data) + padder.finalize()
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=self.backend)
        encryptor = cipher.encryptor()
        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
        return base64.b64encode(encrypted_data).decode('utf-8')

    # AES 解密
    def aes_decrypt(self, encrypted_data, key, iv):
        """
        对AES加密的数据进行解密，返回解密后的明文
        
        Args:
            encrypted_data (str): AES加密后的Base64编码字符串
            key (bytes): AES加密的密钥，长度为16、24或32字节
            iv (bytes): AES加密的初始化向量，长度为16字节
        
        Returns:
            str: 解密后的明文
        
        """
        encrypted_data = base64.b64decode(encrypted_data)
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=self.backend)
        decryptor = cipher.decryptor()
        padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
        unpadder = padding.PKCS7(128).unpadder()
        data = unpadder.update(padded_data) + unpadder.finalize()
        return data.decode('utf-8')  # 将字节转换为字符串

    # 生成 AES 密钥和 IV
    def generate_aes_key_iv(self):
        """
        生成AES加密所需的密钥(key)和初始化向量(iv)。
        
        Args:
            无。
        
        Returns:
            tuple: 一个包含密钥和初始化向量的元组，其中：
                - key (bytes): 长度为32字节的密钥，用于AES加密。
                - iv (bytes): 长度为16字节的初始化向量，用于AES加密。
        
        """
        key = os.urandom(32)  # 256-bit key
        iv = os.urandom(16)   # 128-bit IV
        return key, iv

    # 生成 RSA 密钥对
    def generate_rsa_key_pair(self):
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=self.backend
        )
        public_key = private_key.public_key()
        return private_key, public_key

    # RSA 加密
    def rsa_encrypt(self, data, public_key):
        """
        使用RSA公钥对数据进行加密
        
        Args:
            data (str): 待加密的字符串数据
            public_key (cryptography.hazmat.primitives.asymmetric.rsa.RSAPublicKey): RSA公钥对象
        
        Returns:
            str: 加密后的Base64编码字符串
        
        """
        data = data.encode('utf-8')  # 将字符串转换为字节
        encrypted_data = public_key.encrypt(
            data,
            asym_padding.OAEP(
                mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return base64.b64encode(encrypted_data).decode('utf-8')

    # RSA 解密
    def rsa_decrypt(self, encrypted_data, private_key):
        """
        使用私钥对RSA加密的数据进行解密。
        
        Args:
            encrypted_data (str): 经过RSA加密的Base64编码字符串。
            private_key (cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey): 用于解密的私钥。
        
        Returns:
            str: 解密后的明文数据。
        
        """
        encrypted_data = base64.b64decode(encrypted_data)
        decrypted_data = private_key.decrypt(
            encrypted_data,
            asym_padding.OAEP(
                mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return decrypted_data.decode('utf-8')  # 将字节转换为字符串

    # 计算 SHA-256 哈希
    def sha256_hash(self, data):
        """
        计算给定字符串的SHA256哈希值。
        
        Args:
            data (str): 待计算哈希值的字符串。
        
        Returns:
            str: 计算得到的SHA256哈希值，以十六进制字符串形式返回。
        
        """
        data = data.encode('utf-8')  # 将字符串转换为字节
        hash_object = hashlib.sha256()
        hash_object.update(data)
        hash_digest = hash_object.hexdigest()
        return hash_digest

# 使用示例
# if __name__ == "__main__":
#     crypto_utils = CryptoUtils()
#
#     # AES 示例
#     aes_key, aes_iv = crypto_utils.generate_aes_key_iv()
#     # print(aes_key, aes_iv)
#     aes_data = "Hello, AES!"
#     aes_encrypted = crypto_utils.aes_encrypt(aes_data, aes_key, aes_iv)
#     aes_decrypted = crypto_utils.aes_decrypt(aes_encrypted, aes_key, aes_iv)
#     print(f"AES Encrypted: {aes_encrypted}")
#     print(f"AES Decrypted: {aes_decrypted}")
#
#     # RSA 示例
#     rsa_private_key, rsa_public_key = crypto_utils.generate_rsa_key_pair()
#     rsa_data = "Hello, RSA!"
#     rsa_encrypted = crypto_utils.rsa_encrypt(rsa_data, rsa_public_key)
#     rsa_decrypted = crypto_utils.rsa_decrypt(rsa_encrypted, rsa_private_key)
#     print(f"RSA Encrypted: {rsa_encrypted}")
#     print(f"RSA Decrypted: {rsa_decrypted}")
#
#     # SHA-256 示例
#     sha_data = "Hello, SHA-256!"
#     sha_hash = crypto_utils.sha256_hash(sha_data)
#     print(f"SHA-256 Hash: {sha_hash}")
